
import copy
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from sklearn.externals import joblib

def data_cut(origin_data: pd.DataFrame, target_column: str, test_size: float, random_state=739):
    """
    封装了数据切分的方法,返回测试集训练集的X,y
    from sklearn.model_selection import train_test_split
    ======================================================
    :param origin_data:原始数据
    :param target_column:目标列
    :param test_size:测试集比例
    :param random_state:随机种子
    :return:X_train, X_test, y_train, y_test
    """
    y_data = copy.deepcopy(origin_data[target_column])
    X_data = origin_data.drop(columns=target_column)
    X_train, X_test, y_train, y_test = train_test_split(X_data, y_data, test_size=test_size, random_state=random_state)
    return X_train, X_test, y_train, y_test

def pre_ohe(mid_data: pd.DataFrame, feature_list: list, dump: bool, ohe_path: str):
    """
    :param mid_data: 来源数据
    :param feature_list: 准独热化特征
    :param dump:是否存储
    :param ohe_path: ohe.pkl存储路径
    :return: 独热化数据
    """
    if dump:
        ohe = OneHotEncoder(handle_unknown='ignore')
        ohe.fit(mid_data[feature_list])
        joblib.dump(ohe, ohe_path)
    else:
        ohe = joblib.load(ohe_path)
    mid_data = ohe.transform(mid_data[feature_list].fillna(0)).toarray()
    return mid_data

def pre_labeler(origin: pd.DataFrame, feature_list: list, dump: bool, str2int_path: str):
    """
    :param origin: 原始数据
    :param feature_list: 准标签化特征
    :param dump:调用判断
    :param str2int_path:str2int.pkl存储路径
    :return:标签化数据
    """
    if dump:
        le_name_mapping = {}
        class_le = LabelEncoder()
        for i in range(len(feature_list)):
            class_le.fit(origin[feature_list[i]].values)
            le_name_mapping.update(
                {feature_list[i]: dict(zip(class_le.classes_, class_le.transform(class_le.classes_)))})
            origin[feature_list[i]] = class_le.fit_transform(origin[feature_list[i]].values)  # str2int后的数据
        joblib.dump(le_name_mapping, str2int_path)
        return origin
    else:
        le_name_mapping = joblib.load(str2int_path)
        new_data = pd.DataFrame(columns=feature_list)
        for each in feature_list:
            new_data[each] = origin[each].map(lambda s: le_name_mapping[each].get(s))
        return new_data












